import asyncio
import websockets
import json
import cv2
from yolov5 import YOLOv5

# 加载YOLOv5模型
model = YOLOv5('./yolov5s.pt')  # 请确保路径正确并且模型已经下载

async def detect_objects(video_path, websocket):
    # 使用 OpenCV 打开 RTSP 流
    cap = cv2.VideoCapture(video_path)
    
    if not cap.isOpened():
        await websocket.send(json.dumps({"error": "无法打开视频流"}))
        return
    
    print("RTSP 流已成功打开，开始读取视频帧...")
    
    while True:
        ret, frame = cap.read()
        
        # 检查帧是否读取成功
        if not ret:
            await websocket.send(json.dumps({"error": "无法读取视频帧"}))
            break
        
        print("成功读取视频帧，开始进行目标检测...")
        
        # YOLOv5检测
        results = model.predict(frame)  # 使用 .predict() 方法进行推理
        detections = results.pandas().xywh[0]  # 获取检测结果
        
        if detections.empty:
            await websocket.send(json.dumps({"message": "没有检测到任何对象"}))
            break
        
        # 创建响应
        response = []
        for _, detection in detections.iterrows():
            x_center, y_center, width, height = detection['xcenter'], detection['ycenter'], detection['width'], detection['height']
            xmin = int((x_center - width / 2) * frame.shape[1])
            ymin = int((y_center - height / 2) * frame.shape[0])
            xmax = int((x_center + width / 2) * frame.shape[1])
            ymax = int((y_center + height / 2) * frame.shape[0])

            result = {
                "class": detection['name'],
                "confidence": float(detection['confidence']),
                "coordinates": {
                    "x_min": xmin,
                    "y_min": ymin,
                    "x_max": xmax,
                    "y_max": ymax
                }
            }
            response.append(result)

        # 发送检测结果
        await websocket.send(json.dumps(response))  
    
    cap.release()
    print("视频处理完毕，等待下一个请求...")

# 处理每个WebSocket连接
async def handle_connection(websocket, path):
    try:
        while True:  # 保持连接，等待客户端新的请求
            # 接收客户端发送的消息
            message = await websocket.recv()  # 接收客户端传来的数据
            data = json.loads(message)  # 解析 JSON 数据
            
            # 获取视频路径
            video_path = data.get('video_path')  # 获取视频路径
            if not video_path:
                await websocket.send(json.dumps({"error": "没有提供视频路径"}))
                continue  # 如果没有提供路径，继续等待下一个请求
            
            # 开始处理视频
            await detect_objects(video_path, websocket)

    except Exception as e:
        print(f"Error: {e}")
        await websocket.send(json.dumps({"error": str(e)}))
    finally:
        print("Closing WebSocket connection.")
        await websocket.close()  # 确保连接被正确关闭

# 启动WebSocket服务器
async def main():
    async with websockets.serve(handle_connection, "localhost", 8765):
        print("WebSocket server started at ws://localhost:8765")
        await asyncio.Future()  # 永远运行直到被中断

# 启动服务器
if __name__ == "__main__":
    asyncio.run(main())
